
<!DOCTYPE html>
<html>
  <head>
    
<meta charset="utf-8" >

<title>百度的《AIBox: CTR Prediction Model Training on a Single Node》 | dragon</title>
<meta name="description" content="邮箱(base64)：MTY5MDMwMjk2M0BxcS5jb20=
">

<meta name="viewport" content="width=device-width, initial-scale=1, maximum-scale=1, user-scalable=no">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/animate.css/3.7.0/animate.min.css">

<link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.7.2/css/all.css" integrity="sha384-fnmOCqbTlWIlj8LyTjo7mOUStjsKC4pOpQbqyi7RrhN7udi9RwhKkMHpvLbHG9Sr" crossorigin="anonymous">
<link rel="shortcut icon" href="https://dragonfive.gitee.io//favicon.ico?v=1740893463017">
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.10.0/katex.min.css">
<link rel="stylesheet" href="https://dragonfive.gitee.io//styles/main.css">



<script src="https://cdn.jsdelivr.net/npm/vue/dist/vue.js"></script>
<script src="//cdn.jsdelivr.net/gh/highlightjs/cdn-release@11.5.1/build/highlight.min.js"></script>



  </head>
  <body>
    <div id="app" class="main">
      <div class="site-header-container">
  <div class="site-header">
    <div class="left">
      <a href="https://dragonfive.gitee.io/">
        <img class="avatar" src="https://dragonfive.gitee.io//images/avatar.png?v=1740893463017" alt="" width="32px" height="32px">
      </a>
      <a href="https://dragonfive.gitee.io/">
        <h1 class="site-title">dragon</h1>
      </a>
    </div>
    <div class="right">
      <transition name="fade">
        <i class="icon" :class="{ 'icon-close-outline': menuVisible, 'icon-menu-outline': !menuVisible }" @click="menuVisible = !menuVisible"></i>
      </transition>
    </div>
  </div>
</div>

<transition name="fade">
  <div class="menu-container" style="display: none;" v-show="menuVisible">
    <div class="menu-list">
      
        
          <a href="/" class="menu purple-link">
            首页
          </a>
        
      
        
          <a href="/archives" class="menu purple-link">
            归档
          </a>
        
      
        
          <a href="/tags" class="menu purple-link">
            标签
          </a>
        
      
        
          <a href="/post/about" class="menu purple-link">
            关于
          </a>
        
      
    </div>
  </div>
</transition>


      <div class="content-container">
        <div class="post-detail">
          
          <h2 class="post-title">百度的《AIBox: CTR Prediction Model Training on a Single Node》</h2>
          <div class="post-info post-detail-info">
            <span><i class="icon-calendar-outline"></i> 2020-12-11</span>
            
              <span>
                <i class="icon-pricetags-outline"></i>
                
                  <a href="https://dragonfive.gitee.io/tag/xun-lian-kuang-jia/">
                    训练框架
                    
                      ，
                    
                  </a>
                
                  <a href="https://dragonfive.gitee.io/tag/gd01ID-F3/">
                    论文阅读
                    
                      ，
                    
                  </a>
                
                  <a href="https://dragonfive.gitee.io/tag/WzibKNMac/">
                    深度学习
                    
                  </a>
                
              </span>
            
          </div>
          <div class="post-content" v-pre>
            <p>AIBox 是百度提出的训练框架，论文 AIBox: CTR Prediction Model Training on a Single Node 进行了相关介绍。</p>
<h1 id="一-aibox-的技术创新及优势">一、AIBox 的技术创新及优势</h1>
<p>创新点与动机：<br>
AIBox 的核心想法就是想在一台机器上用GPU加速训练，但是参数实在太大了，所以就把计算密集的模型运算部分（joint learning）放在GPU完成，把取embedding 部分放在cpu部分完成(embedding learning)，这就是AiBox 的第一个创新：把网络切分为两部分。</p>
<p>但即便主存用了1TB 的内存，embedding 还是太大了，10^12 个key，每个key 的 weight 即使用8个字节存，key用8个字节存，也需要1.6TB , 所以论文提出了第二个创新：把embedding 存在SSD上，同时为了降低延迟和减少写操作对ssd寿命的影响，建立了二级缓存机制。</p>
<p>为了提高速度，AIBOX使用了流水线，把从hdfs 读数据(socket IO)，从SSD查Embedding （SSD io） 和 cpu+gpu 计算组成3阶段的 pipeline</p>
<p>优势：<br>
AIBOX不存在像分布式系统普遍存在的网络通信开销问题，然后在系统稳定性方面AIBOX与具有数千台机器的分布式集群相比更加稳定不会轻易宕机，而且在同步开销方面AIBOX只是涉及到一些内存锁和GPU片之间的少量通信。</p>
<h1 id="二-关于网络结构切分">二、关于网络结构切分</h1>
<p>the first module focuses on the embedding learning with high-dimensional &amp; sparse features and the second module is for joint learning with dense features resulted from the first module.</p>
<p>The embedding learning is processed on CPUs to help learn low dimensional dense embedding representations.</p>
<p>By transferring the learned embedding vectors from CPUs to GPUs, the computation-intensive joint learning module can make full use of the powerful GPUs for CTR prediction.</p>
<p>CPU 部分把数据从稀疏特征转化成 embedding （embedding learning），然后把embedding 传到 GPU，GPU进行一轮训练 (joint learning)</p>
<p>论文这部分讲了一些网络的设计细节，但这块感觉跟 AIBox本身没什么关系，论文写到：</p>
<p>把第一隐含层和最后一层隐含层的结果合并起来，第一层包含了low-level 的与输入信息最相关的feature，最后一层包含了high-level 的最抽象和有用的信息。这样会得到更准确的CTR预估结果。</p>
<p>训练两阶段(cpu+gpu)，梯度更新也是两阶段(gpu+cpu)。</p>
<h1 id="三-aibox-架构划分">三、AIBox 架构划分</h1>
<p><strong>架构：</strong><br>
分为三部分：CPU、GPU和 sparse table buff</p>
<ul>
<li>
<p>cpu模块：协调调度和embedding学习<br>
从hdfs读数据(一个pass)，向Sparse Table模块查embedding，然后发给GPU<br>
拿到gpu传来的梯度，更新sparse table<br>
定期save ckpt 到 hdfs</p>
</li>
<li>
<p>sparse table：把10^12 的离散特征的数据存储到ssd上的kv系统里<br>
内存中的key hash 索引存了特征到文件的映射关系，<br>
in-memory cache strategy 构造cache 和 buffer 来减少延迟</p>
</li>
<li>
<p>gpu模块：联合学习<br>
cpu传来的 embedding 被放入 HBMs 中，然后被fed 给 dense 联合学习网络<br>
emb通过pci-e总线进行传输<br>
一个CUDA stream进行数据传输，另一个cuda stream 进行学习计算<br>
HBMs如同片上ps一样工作，<br>
每个pass 在每个gpu 上计算新的参数，各gpu通过NVLink进行同步</p>
</li>
</ul>
<p><strong>3阶段流水线：network, SSDs and CPUs + GPUs</strong></p>
<figure data-type="image" tabindex="1"><img src="https://dragonfive.gitee.io//post-images/1625143570973.png" alt="" loading="lazy"></figure>
<figure data-type="image" tabindex="2"><img src="https://dragonfive.gitee.io//post-images/1625143575338.png" alt="" loading="lazy"></figure>
<h1 id="四-sparsetable-架构">四、sparseTable 架构</h1>
<p>由两部分构成：key hash index and bi-level cache management</p>
<h2 id="一key-hash-index">（一）key hash index</h2>
<p>Key Hash Index 存的是 10^12 个 feature key 到 ssd 文件的映射关系，直接每个key 存一个文件需要1.6TB大小的内存，放不下。</p>
<p>通过对key 取模进行分组建立 group 与 file 的对应关系，放在内存中。</p>
<p>group(key) → key mod 1012/m. We set m = ⌊BLOCK/(8 + sizeof(value))⌋,其中Block 是每次从ssd取数据的最小单元。</p>
<p>hash函数可以通过预训练一个模型来最大化feature 共现，把共现的feature 分到同样的桶里面。</p>
<h2 id="二二级缓存机制">（二）二级缓存机制</h2>
<p>ssd 的延迟是内存的1000倍，ssd是微秒级别延迟，内存是纳秒级别延迟</p>
<p>在一个 pass of mini batch 中只有1%的参数会被用到，所以我们可以用in-memory cache 来存储高频访问的hot parameters</p>
<p>SSD有物理性能限制：每个存储单元只能被写入（擦除）数千次，cache机制可以作为参数缓存，来减少更新参数对SSD使用寿命的影响</p>
<p>使用两个分离的链表进行拉链来提升探测性能。对每个ssd文件使用Bloom filter来减少不必要的读取。</p>
<p><strong>第一级缓存</strong></p>
<p>使用 si =hash1(g_id) 来算出一个 cache slot 槽，对应一个ssd 文件，对于参数并未进行真正初始化，而是在第一次访问到参数的时候，先用 bloom filter 探测key 是否在 slot 集合里，如果不在就不用读取这个文件，而是直接使用默认值，以此来减少不必要的ssd读取。</p>
<p><strong>二级缓存</strong></p>
<p>hash2(g_id, bucket)</p>
<p>对 一级的槽进行分桶bucket，来使得拉的链比较短。bucket 参数通过调节可以权衡空间和探测效率</p>
<p><strong>两条拉链</strong></p>
<ul>
<li>LRU 链用于保存最近访问过的key，以此来减少探测次数</li>
<li>LFU链按访问频次来保存key，用于缓存管理，只有当LFU满了需要删除低频key时，相应的数据才会写回到ssd上面</li>
</ul>
<p>由于经常有链条中的节点进行增删，所以使用线程池以Slab memory allocation mechanism机制 进行管理。</p>
<figure data-type="image" tabindex="3"><img src="https://dragonfive.gitee.io//post-images/1625143705399.png" alt="" loading="lazy"></figure>
<p><strong>文件管理系统</strong><br>
batch产生的小文件对于先有的文件系统有很大的压力，把许多小文件组成一个组来创建较少的文件。小文件的名字由大文件的名字加上offset构成，保存在第一级cache slot 中</p>
<p>监控文件系统的大小，合并访问量少的小问题，当model_size 达到最大冗余度的时候删掉访问少的文件，MAX_REPLICATION=SSD capacity ∗ (85% + overprovisioning)/model size<br>
<img src="https://dragonfive.gitee.io//post-images/1625143710125.png" alt="" loading="lazy"></p>
<h1 id="五-实验部分">五、实验部分</h1>
<p><strong>实验</strong><br>
AIBox 8 个GPU， 服务器级别的cpu, 1T 内存，Raid 0 nvme ssd<br>
MPI集群方式用75个计算节点</p>
<ul>
<li>
<p>AIBox 的硬件和维护费用比集群训练方式少 10%，执行时间多25%</p>
</li>
<li>
<p>AIBox 的auc 比集群方式稍好，可能是因为AIBox 这种单节点的方式，同步参数频率更高</p>
</li>
</ul>
<figure data-type="image" tabindex="4"><img src="https://dragonfive.gitee.io//post-images/1625143749656.png" alt="" loading="lazy"></figure>
<p>六、总结与一些细节问题：<br>
论文介绍了 AIBox 架构的一些细节方面，借助一系列系统设计方案如缓存机制来解决问题，通过这些并不是很fashion的技术合并，论文实现了集中式训练的技术突破，这种通过技术积累然后撬动难题的解决问题的方式值得我们学习。</p>
<p>但是还有一些细节没有讲清楚：</p>
<ol>
<li>
<p>AIBox 有几个worker 进行工作，他们是数据并行，还是使用同样的数据进行训练（文中提到AIBox 会在每个pass of mini batch 进行同步，所以应该不是一个worker 在参与训练）</p>
</li>
<li>
<p>AIBox 使用集中的训练方式，那如果这台机器挂掉，是不是根本没有办法进行恢复，只能另找一个机器从 ckpt 训练</p>
</li>
<li>
<p>文章没有介绍使用的具体计算引擎 (怀疑跟 horovod 接近）</p>
</li>
<li>
<p>同样文章没有介绍参数同步的细节，没有相关 all_reduce 的介绍（可以是使用了一个开源的框架，而这部分论文没有进行改进，所以没有做深入介绍）</p>
</li>
<li>
<p>文章开头提到使用 in-HBM ps 来减少数据传输，但是后面没有详细进行介绍</p>
</li>
</ol>
<p>总体上感觉这篇论文实用性强，但是细节介绍得不多</p>

          </div>
        </div>

        
          <div class="next-post">
            <a class="purple-link" href="https://dragonfive.gitee.io/post/seastar-jiao-cheng-fan-yi/">
              <h3 class="post-title">
                下一篇：seastar教程翻译
              </h3>
            </a>
          </div>
          
      </div>

      

      <div class="site-footer">
  <div class="slogan">邮箱(base64)：MTY5MDMwMjk2M0BxcS5jb20=
</div>
  <div class="social-container">
    
      
        <a href="https://github.com/DragonFive" target="_blank">
          <i class="fab fa-github"></i>
        </a>
      
    
      
    
      
    
      
    
      
    
  </div>
  Powered by <a href="https://github.com/getgridea/gridea" target="_blank">Gridea</a> | <a class="rss" href="https://dragonfive.gitee.io//atom.xml" target="_blank">RSS</a>
</div>


    </div>
    <script type="application/javascript">

hljs.initHighlightingOnLoad()

var app = new Vue({
  el: '#app',
  data: {
    menuVisible: false,
  },
})

</script>




  </body>
</html>
